{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/customization/prompts/chat_prompts.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat Prompts Customization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Setup\n",
    "\n",
    "Lets customize them to always answer, even if the context is not helpful!\n",
    "\n",
    "Using `RichPromptTemplate`, we can define Jinja-formatted prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.prompts import RichPromptTemplate\n",
    "\n",
    "chat_text_qa_prompt_str = \"\"\"\n",
    "{% chat role=\"system\" %}\n",
    "Always answer the question, even if the context isn't helpful.\n",
    "{% endchat %}\n",
    "\n",
    "{% chat role=\"user\" %}\n",
    "The following is some retrieved context:\n",
    "\n",
    "<context>\n",
    "{{ context_str }}\n",
    "</context>\n",
    "\n",
    "Using the context, answer the provided question:\n",
    "{{ query_str }}\n",
    "{% endchat %}\n",
    "\"\"\"\n",
    "text_qa_template = RichPromptTemplate(chat_text_qa_prompt_str)\n",
    "\n",
    "# Refine Prompt\n",
    "chat_refine_prompt_str = \"\"\"\n",
    "{% chat role=\"system\" %}\n",
    "Always answer the question, even if the context isn't helpful.\n",
    "{% endchat %}\n",
    "\n",
    "{% chat role=\"user\" %}\n",
    "The following is some new retrieved context:\n",
    "\n",
    "<context>\n",
    "{{ context_msg }}\n",
    "</context>\n",
    "\n",
    "And here is an existing answer to the query:\n",
    "<existing_answer>\n",
    "{{ existing_answer }}\n",
    "</existing_answer>\n",
    "\n",
    "Using both the new retrieved context and the existing answer, either update or repeat the existing answer to this query:\n",
    "{{ query_str }}\n",
    "{% endchat %}\n",
    "\"\"\"\n",
    "refine_template = RichPromptTemplate(chat_refine_prompt_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Prompts\n",
    "\n",
    "Now, we use the prompts in an index query!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-4o-mini\")\n",
    "Settings.embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-small\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents)\n",
    "\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Before Customizing Templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets see the default existing prompts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'response_synthesizer:text_qa_template': SelectorPromptTemplate(metadata={'prompt_type': <PromptType.QUESTION_ANSWER: 'text_qa'>}, template_vars=['context_str', 'query_str'], kwargs={}, output_parser=None, template_var_mappings={}, function_mappings={}, default_template=PromptTemplate(metadata={'prompt_type': <PromptType.QUESTION_ANSWER: 'text_qa'>}, template_vars=['context_str', 'query_str'], kwargs={}, output_parser=None, template_var_mappings=None, function_mappings=None, template='Context information is below.\\n---------------------\\n{context_str}\\n---------------------\\nGiven the context information and not prior knowledge, answer the query.\\nQuery: {query_str}\\nAnswer: '), conditionals=[(<function is_chat_model at 0x1062d16c0>, ChatPromptTemplate(metadata={'prompt_type': <PromptType.CUSTOM: 'custom'>}, template_vars=['context_str', 'query_str'], kwargs={}, output_parser=None, template_var_mappings=None, function_mappings=None, message_templates=[ChatMessage(role=<MessageRole.SYSTEM: 'system'>, additional_kwargs={}, blocks=[TextBlock(block_type='text', text=\"You are an expert Q&A system that is trusted around the world.\\nAlways answer the query using the provided context information, and not prior knowledge.\\nSome rules to follow:\\n1. Never directly reference the given context in your answer.\\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\")]), ChatMessage(role=<MessageRole.USER: 'user'>, additional_kwargs={}, blocks=[TextBlock(block_type='text', text='Context information is below.\\n---------------------\\n{context_str}\\n---------------------\\nGiven the context information and not prior knowledge, answer the query.\\nQuery: {query_str}\\nAnswer: ')])]))]),\n",
       " 'response_synthesizer:refine_template': SelectorPromptTemplate(metadata={'prompt_type': <PromptType.REFINE: 'refine'>}, template_vars=['query_str', 'existing_answer', 'context_msg'], kwargs={}, output_parser=None, template_var_mappings={}, function_mappings={}, default_template=PromptTemplate(metadata={'prompt_type': <PromptType.REFINE: 'refine'>}, template_vars=['query_str', 'existing_answer', 'context_msg'], kwargs={}, output_parser=None, template_var_mappings=None, function_mappings=None, template=\"The original query is as follows: {query_str}\\nWe have provided an existing answer: {existing_answer}\\nWe have the opportunity to refine the existing answer (only if needed) with some more context below.\\n------------\\n{context_msg}\\n------------\\nGiven the new context, refine the original answer to better answer the query. If the context isn't useful, return the original answer.\\nRefined Answer: \"), conditionals=[(<function is_chat_model at 0x1062d16c0>, ChatPromptTemplate(metadata={'prompt_type': <PromptType.CUSTOM: 'custom'>}, template_vars=['context_msg', 'query_str', 'existing_answer'], kwargs={}, output_parser=None, template_var_mappings=None, function_mappings=None, message_templates=[ChatMessage(role=<MessageRole.USER: 'user'>, additional_kwargs={}, blocks=[TextBlock(block_type='text', text=\"You are an expert Q&A system that strictly operates in two modes when refining existing answers:\\n1. **Rewrite** an original answer using the new context.\\n2. **Repeat** the original answer if the new context isn't useful.\\nNever reference the original answer or context directly in your answer.\\nWhen in doubt, just repeat the original answer.\\nNew Context: {context_msg}\\nQuery: {query_str}\\nOriginal Answer: {existing_answer}\\nNew Answer: \")])]))])}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_engine.get_prompts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And how do they respond when asking about unrelated concepts?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The provided information does not contain any details about Joe Biden.\n"
     ]
    }
   ],
   "source": [
    "print(query_engine.query(\"Who is Joe Biden?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### After Customizing Templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can update the templates and observe the change in response!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine.update_prompts(\n",
    "    {\n",
    "        \"response_synthesizer:text_qa_template\": text_qa_template,\n",
    "        \"response_synthesizer:refine_template\": refine_template,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Joe Biden is the 46th President of the United States, having taken office on January 20, 2021. He previously served as Vice President under President Barack Obama from 2009 to 2017 and has had a long career in politics, including serving as a U.S. Senator from Delaware from 1973 to 2009. Biden is a member of the Democratic Party and has focused on issues such as healthcare, climate change, and economic recovery during his presidency.\n"
     ]
    }
   ],
   "source": [
    "print(query_engine.query(\"Who is Joe Biden?\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
